
import os
from logging import warning

import numpy as np
from download import download
import matplotlib.pyplot as plt
from mindspore.dataset import vision
from mindspore.dataset import MnistDataset, GeneratorDataset


def download_dataset(url, path="./"):
    """
    通过 download 下载数据集
    :param url: 下载链接
    :param path: 数据集保存地址
    :return:
    """
    try:
        if os.path.exists(path):
            print("{} 文件以存在".format(path))
        else:
            path = download(url, path, kind="zip", replace=True)
            print("下载完成：{}".format(path))

    except RuntimeWarning as e:
        print("数据集下载失败：{}".format(e))





def show_visualize(dataset):
    # 创建一个画布
    figure = plt.figure(figsize=(4, 4))
    cols, rows = 3, 3

    plt.subplots_adjust(wspace=0.5, hspace=0.5)

    for idx, (image, label) in enumerate(dataset.create_tuple_iterator()):
        figure.add_subplot(rows, cols, idx + 1)
        plt.title(label)
        plt.axis("off")
        plt.imshow(image.asnumpy().squeeze(), cmap="gray")
        if idx == cols * rows - 1:
            break

    plt.show()


def show_shuffle(data_path='./mnist/MNIST_Data/train'):
    dataset = MnistDataset(data_path, shuffle=False)
    dataset = dataset.shuffle(buffer_size=64)
    show_visualize(dataset)


def show_map(data_path='./mnist/MNIST_Data/train'):
    dataset = MnistDataset(data_path, shuffle=False)
    image, label = next(dataset.create_tuple_iterator())
    print("数据的列名")
    print(dataset.create_dict_iterator().get_col_names())

    print("数据类型调整前:")
    print(image.shape, image.dtype)

    print("数据类型调整后:")
    dataset = dataset.map(vision.Rescale(1.0 / 255.0, 0), input_columns='image')
    image, label = next(dataset.create_tuple_iterator())
    print(image.shape, image.dtype)


def show_batch(data_path='./mnist/MNIST_Data/train'):
    dataset = MnistDataset(data_path, shuffle=False)
    dataset_32 = dataset.batch(batch_size=32)
    image, label = next(dataset_32.create_tuple_iterator())
    print("batch 为 32 时,每次迭代获取的样例:")
    print(image.shape, image.dtype)

    dataset = MnistDataset(data_path, shuffle=False)
    dataset_128 = dataset.batch(batch_size=128)
    image, label = next(dataset_128.create_tuple_iterator())
    print("batch 为 128 时,每次迭代获取的样例:")
    print(image.shape, image.dtype)


class RandomAccessDataset:
    def __init__(self):
        self._data = np.ones((5, 2))
        self._label = np.zeros((5, 1))

    def __getitem__(self, index):
        return self._data[index], self._label[index]

    def __len__(self):
        return len(self._data)


def show_getitem_dataset():
    loader = RandomAccessDataset()

    dataset = GeneratorDataset(source=loader, column_names=["data", "label"])

    for data in dataset:
        print(data)


class IterableDataset:
    def __init__(self):
        self._data = np.ones((5, 2))
        self._label = np.zeros((5, 1))
        self._index = len(self._label) + 1

    def __next__(self):
        if next(self.index):
            print(self.index)
            return next(self.data), next(self.label)

    def __iter__(self):
        # self.index = iter(range(1, self._index))
        self.index = iter(self.breaker, 3)
        self.data = iter(self._data)
        self.label = iter(self._label)
        return self

    def breaker(self):
        self._index -= 1
        return self._index


def show_iter_dataset():
    loader = IterableDataset()

    dataset = GeneratorDataset(source=loader, column_names=["data", "label"])

    for data in dataset:
        print(data)




if __name__ == '__main__':
    # # 展示数据集
    # url_mnist = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/MNIST_Data.zip"
    # download_dataset(url_mnist, './mnist')
    # train_dataset = MnistDataset('./mnist/MNIST_Data/train', shuffle=False)
    # show_visualize(train_dataset)

    # # 展示 shuffle
    # show_shuffle(data_path='./mnist/MNIST_Data/train')
    # show_shuffle(data_path='./mnist/MNIST_Data/train')

    # show_map()

    # show_batch()

    # # 展示数据集
    # show_dataset()

    # 展示数据集
    show_iter_dataset()
    # index = iter(range(1, 5), 3)
    # for i in index:
    #     print(i)